package com.ruoyi.common.utils;


import com.ruoyi.common.core.CalcThreadManager;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;


@Slf4j
@Component
public class SseEmitterUtils {

    // 当前连接数
    private static AtomicInteger count = new AtomicInteger(0);
    // 存储 SseEmitter 信息
    private static volatile Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();

    //用来存储每个通道最后一次发送的信息，在用户切回页面时，返回。
    private static volatile Map<String, String> msgCacheMap = new ConcurrentHashMap<>();

    /**
     * 链接，不返回任何信息
     * @param key
     * @param flag
     * @return
     */
    public static SseEmitter connectWithNoMessage(String key,String flag) {
        if (sseEmitterMap.containsKey(key)) {
            log.debug("sse-connect时，sseEmitterMap已存在key:{} , 缓存信息:{}",key, msgCacheMap.get(key));
            SseEmitter sseEmitter = sseEmitterMap.get(key);
            String message = msgCacheMap.get(key);
            if ("0".equals(flag) && message != null){
                CompletableFuture.runAsync(() -> {
                    try {
                        TimeUnit.MILLISECONDS.sleep(400);
                        sendMessage(key, message);
                    } catch (InterruptedException e) {
                        throw new RuntimeException(e);
                    }
                });
            }
            return sseEmitter;
        }
        try {
            // 设置超时时间，0表示不过期。默认30秒
            SseEmitter sseEmitter = new SseEmitter(0L);
            // 注册回调
            sseEmitter.onCompletion(completionCallBack(key));
            sseEmitter.onError(errorCallBack(key));
            sseEmitter.onTimeout(timeoutCallBack(key));
            sseEmitterMap.put(key, sseEmitter);
            // 数量+1
            count.getAndIncrement();
            if (flag.equals("1")){
                log.info("sse-connect创建新的连接，key:{}",key);
            }else{
                sseEmitter.send("200");
                remove(key);
            }
            return sseEmitter;
        } catch (Exception e) {
            log.info("创建新的SSE连接异常，当前连接Key为：{}", key);
        }
        return null;
    }

    /**
     * 创建用户连接并返回 SseEmitter
     * @param key
     * @return
     */
    public static SseEmitter connect(String key,String flag) {
        if (sseEmitterMap.containsKey(key)) {
            log.debug("sse-connect时，sseEmitterMap已存在key:{} , 缓存信息:{}",key, msgCacheMap.get(key));
            SseEmitter sseEmitter = sseEmitterMap.get(key);
            String message = msgCacheMap.get(key);
            if ("0".equals(flag) && message != null){
                CompletableFuture.runAsync(() -> {
                    try {
                        TimeUnit.MILLISECONDS.sleep(400);
                        sendMessage(key, message);
                    } catch (InterruptedException e) {
                        throw new RuntimeException(e);
                    }
                });
            } else if ("1".equals(flag)) {
                CompletableFuture.runAsync(() -> {
                    try {
                        TimeUnit.MILLISECONDS.sleep(400);
                        sendMessage(key, "0");
                    } catch (InterruptedException e) {
                        throw new RuntimeException(e);
                    }
                });
            }
            return sseEmitter;
        }
        try {
            // 设置超时时间，0表示不过期。默认30秒
            SseEmitter sseEmitter = new SseEmitter(0L);
            // 注册回调
            sseEmitter.onCompletion(completionCallBack(key));
            sseEmitter.onError(errorCallBack(key));
            sseEmitter.onTimeout(timeoutCallBack(key));
            sseEmitterMap.put(key, sseEmitter);
            // 数量+1
            count.getAndIncrement();
            if (flag.equals("1")){
                sseEmitter.send("0");
                log.debug("sse-connect创建新的连接，key:{}",key);
            }else{
                sseEmitter.send("200");
                remove(key);
            }
            return sseEmitter;
        } catch (Exception e) {
            log.info("创建新的SSE连接异常，当前连接Key为：{}", key);
        }
        return null;
    }

    /**
     * 给指定用户发送消息
     **/
    public static void sendMessage(String key, String message) {
        if(!StringUtils.hasLength(key)) {
            return;
        }
        if (sseEmitterMap.containsKey(key)) {
            try {
                log.debug("准备发送消息，key:{}, message:{}",key, message);
                sseEmitterMap.get(key).send(message);
                msgCacheMap.put(key,message);
            } catch (Exception e) {
                log.error("用户[{}]推送异常, message:{}, 异常：{}", key, message, e.getMessage());
                //remove(key);
            }
        }else {
            log.error("发送消息失败，key:{} 不存在, message:{}",key, message);
        }
    }

    /**
     * 移除连接
     **/
    public static void remove(String key) {
        if (sseEmitterMap.containsKey(key)){

            try {
                SseEmitter sseEmitter = sseEmitterMap.get(key);
                sseEmitter.complete();
            }catch (Exception e){
                log.info("发送complete失败, key:{}", key);
            }
            sseEmitterMap.remove(key);
            msgCacheMap.remove(key);
            // 数量-1
            count.getAndDecrement();
            //移除时,也将缓存线程清空
            CalcThreadManager.clearProcess(key);
        }

        log.info("已移除连接：{}", key);
    }

    /**
     *
     * 获取当前连接信息
     **/
    public static List<String> getIds() {
        return new ArrayList<>(sseEmitterMap.keySet());
    }

    /**
     * 获取当前连接数量
     **/
    public static int getCount() {
        return count.intValue();
    }

    private static Runnable completionCallBack(String key) {
        return () -> {
            log.info("完成连接：{}", key);
            remove(key);
        };
    }

    private static Runnable timeoutCallBack(String key) {
        return () -> {
            log.info("连接超时：{}", key);
            remove(key);
        };
    }

    private static Consumer<Throwable> errorCallBack(String key) {
        return throwable -> {
            log.info("连接异常：{}", key);
            remove(key);
        };
    }

    public static void removeByUser(String username) {
        log.info("正在移除用户:{}，所有的sse连接", username);
        //判断该用户下是否存在已开启通道，存在则进行关闭
        Iterator<Map.Entry<String, SseEmitter>> iterator = sseEmitterMap.entrySet().iterator();
        while (iterator.hasNext()) {
            Map.Entry<String, SseEmitter> entry = iterator.next();
            if (entry.getKey().contains(username)) {
                log.info("正在移除用户:{}，sse连接key:{}", username, entry.getKey());
                iterator.remove();
                msgCacheMap.remove(entry.getKey());
            }
        }
    }

    public static Map<String, SseEmitter> getSseEmitterMap() {
        return sseEmitterMap;
    }
}
